import config
from pathlib import Path
from obs import ObsClient
from groundingdino.util.inference import load_model
from mine import mine, data_path_dict
import logging
import logging.handlers
import sys
logger = logging.getLogger("my_log")
logger.setLevel(logging.INFO)
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
file_handler = logging.handlers.RotatingFileHandler("my_log.log", maxBytes=10**6, backupCount=3)
file_handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)
logger.addHandler(console_handler)
logger.addHandler(file_handler)

if __name__ == '__main__':
    try:
        #in input param, there is datasetId, datasetType, datasetPath, imageConfTh, pcdConfTh
        if len(sys.argv) != 10:
            error_message = "parameter number is {}, need 10 including data_mining.py".format(len(sys.argv))
            logger.error(error_message)
            raise Exception(error_message)
        taskId = sys.argv[1]
        datasetPath = sys.argv[2]
        startFrameNo = int(sys.argv[3])
        endFrameNo = int(sys.argv[4])
        labelsId = sys.argv[5]
        prompts = sys.argv[6]
        boxTh = float(sys.argv[7])
        textTh = float(sys.argv[8])
        callbackUrl = sys.argv[9]
        logger.info("data mining container receive data:{}".format(sys.argv))
        # check param
        if not isinstance(datasetPath, str):
            error_message = "datasetPath {} type check failed, need str".format(datasetPath)
            logger.error(error_message)
            raise Exception(error_message)

        if not datasetPath.endswith("/"):
            datasetPath = datasetPath + "/" 
        
        if not isinstance(callbackUrl, str):
            error_message = "callbackUrl {} type check failed, need str".format(callbackUrl)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(labelsId, str):
            error_message = "labelsId {} type check failed, need str".format(labelsId)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(prompts, str):
            error_message = "prompts {} type check failed, need str".format(prompts)
            logger.error(error_message)
            raise Exception(error_message)
        
        if "." not in prompts:
            error_message = "prompts {} value check failed, must exist at least one '.'".format(prompts)
            logger.error(error_message)
            raise Exception(error_message)
        
        labelsId_len = len(labelsId.split(","))
        prompts_len = len(prompts.split(".")) - 1 #"cat . dog . cow ."
        if not labelsId_len == prompts_len:
            error_message = "prompts {} len:{} and labelsId {} len:{} check failed, need to be same".format(prompts, prompts_len, labelsId, labelsId_len)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(startFrameNo, int):
            error_message = "startFrameNo {} type check failed, need int".format(startFrameNo)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(endFrameNo, int):
            error_message = "endFrameNo {} type check failed, need int".format(endFrameNo)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not endFrameNo >= startFrameNo:
            error_message = "endFrameNo {} and startFrameNo {} check failed, endFrameNo should be bigger".format(endFrameNo, startFrameNo)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(boxTh, float):
            error_message = "boxTh {} type check failed, need float".format(boxTh)
            logger.error(error_message)
            raise Exception(error_message)
        
        if boxTh < 0.000001 or boxTh > 0.999999: 
            error_message = "boxTh {} value check failed, need [0.01, 0.99]".format(boxTh)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(textTh, float):
            error_message = "textTh {} type check failed, need float".format(textTh)
            logger.error(error_message)
            raise Exception(error_message)
        
        if textTh < 0.00001 or textTh > 0.999999: 
            error_message = "textTh {} value check failed, need [0.01, 0.99]".format(textTh)
            logger.error(error_message)
            raise Exception(error_message)
        logger.info("data mining container receive data checked ok")
        
        # check bucket download
        obs_client_annotation = ObsClient(access_key_id=config.annotation_ak, secret_access_key=config.annotation_sk, server=config.annotation_endpoint)
        obs_client_datacenter = ObsClient(access_key_id=config.ak, secret_access_key=config.sk, server=config.endpoint)
        logger.info("obs_client init done")
        response = obs_client_datacenter.listObjects(config.bucket_name, prefix=datasetPath)#
        if response.status < 300:
            logger.info('{} listObjects success'.format(config.bucket_name + "/" + datasetPath))
        else:
            error_message = '{} listObjects failed'.format(config.bucket_name + "/" + datasetPath)
            logger.error(error_message)
            raise Exception(error_message)
        
        #check zip file
        objects = []
        for obj in response.body.contents:
            if obj.key.endswith(".zip"):
                objects.append(obj.key)
        if len(objects) == 0:
            error_message = 'no zip file in {}'.format(config.bucket_name + "/" + datasetPath)
            logger.error(error_message)
            raise Exception(error_message)
        else:
            logger.info('{} zip files in {}'.format(len(objects), config.bucket_name + "/" + datasetPath))

        # check bucket upload 
        object_key = data_path_dict["mining_upload_path"] + datasetPath
        response = obs_client_datacenter.putFile(config.bucket_name, object_key+"test_upload.json", file_path="./test_upload.json")
        if response.status < 300:
            logger.info('{} putFile test_upload success'.format(config.bucket_name + "/" + object_key))
        else:
            error_message = '{} putFile test_upload failed'.format(config.bucket_name + "/" + object_key)
            logger.error(error_message)
            raise Exception(error_message)
        
        # start mine
        groundingdino_model = load_model("./img_detection_groundingdino/groundingdino/config/GroundingDINO_SwinT_OGC.py", "./img_detection_groundingdino/weights/groundingdino_swint_ogc.pth")
        groundingdino_model.cuda()
        groundingdino_model.eval()
        logger.info("groundingdino model load done")
        mine(groundingdino_model, obs_client_datacenter, sys.argv, logger)
        logger.info("data mining function done")

    except Exception as e:
        logger.error("/dataMining meet error {}".format(e))

    
